import codecs
import os
from hexbytes import HexBytes
from Crypto.Cipher import AES
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes


class CryptoTest(object):

    @classmethod
    def fernet_1(cls):
        # Put this somewhere safe!
        key = Fernet.generate_key()
        f = Fernet(key)
        s = "典12s!@+==s;l、‘’ll][]lkvk，lk"
        # token = f.encrypt(b"A really secret message. Not for prying eyes.")
        token = f.encrypt(s.encode(encoding='utf8'))
        print(token)

        res = f.decrypt(token)
        print(res, type(res))
        print(res.decode())

    @classmethod
    def aes_1(cls):
        password = 'dahsjhfjakhfjksdsssss1sx'
        iv = 'fdshgjdsngjknaaa'
        print(len(password), len(iv))
        ##这边填充方式没有自带的方法,需要自己写
        pad = lambda s: s + (16 - len(s) % 16) * chr(0)  # ZeroPadding
        # pad = lambda s: s + (16 - len(s) % 16) * chr(16 - len(s) % 16)  # Pkcs7
        obj = AES.new(password.encode(), AES.MODE_CBC, iv.encode())
        encrypted = obj.encrypt(pad("helloworld").encode("utf8"))
        print(encrypted)
        # b'k\xb8\xad\x1c\xde\xa3O1n-\xf9\x99\xbc\x16\xb4\xd7'

        obj2 = AES.new(password.encode(), AES.MODE_CBC, iv.encode())
        print(obj2.decrypt(encrypted).decode())  # 结果也需要手动去掉填充...
        # b'helloworld\x06\x06\x06\x06\x06\x06'

    @classmethod
    def aes_2(cls):
        # 来获取一个指定长度的bytes对象, 它实际上是在获取不同操作系统特定提供的随机源, 它可以被用来做随机加密的key使用
        key = os.urandom(32)
        iv = os.urandom(16)
        print(key, len(key))
        print(iv, len(iv))
        cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
        encryptor = cipher.encryptor()
        # the buffer needs to be at least len(data) + n - 1 where n is cipher/mode block size in bytes
        buf = bytearray(31)
        print(buf)
        print(len("a secret message"))
        len_encrypted = encryptor.update_into(b"a secret message", buf)
        print(len_encrypted)
        # get the ciphertext from the buffer reading only the bytes written to it (len_encrypted)
        ct = bytes(buf[:len_encrypted]) + encryptor.finalize()
        print(ct)

        # 解密
        decryptor = cipher.decryptor()
        len_decrypted = decryptor.update_into(ct, buf)
        # get the plaintext from the buffer reading only the bytes written (len_decrypted)
        print(bytes(buf[:len_decrypted]) + decryptor.finalize())

    @classmethod
    def rsa(cls):
        private_key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=2048,
        )
        print(private_key)
        pem = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.BestAvailableEncryption(b'mypassword')
        )
        print(pem)

        pem = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.TraditionalOpenSSL,
            encryption_algorithm=serialization.NoEncryption()
        )

        print(pem)

        # 私钥加签
        message = b"A message I want to sign"
        signature = private_key.sign(
            message,
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()),
                salt_length=padding.PSS.MAX_LENGTH
            ),
            hashes.SHA256()
        )
        print(signature)

        # 公钥验签
        public_key = private_key.public_key()
        s = public_key.verify(
            signature,
            message,
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()),
                salt_length=padding.PSS.MAX_LENGTH
            ),
            hashes.SHA256()
        )

        # 公钥加密
        message = b"encrypted data"
        ciphertext = public_key.encrypt(
            message,
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        )
        print("公钥加密:", ciphertext)
        het_str = codecs.decode(codecs.encode(HexBytes(ciphertext), 'hex'), 'ascii')
        print(het_str)
        print(HexBytes(f"0x{het_str}"))

        # 私钥解密
        plaintext = private_key.decrypt(
            ciphertext,
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        )
        print(plaintext)


if __name__ == '__main__':
    CryptoTest.rsa()
